package cn.hy.core;

import cn.hy.model.LuaScript;
import cn.hy.model.Result;
import cn.hy.model.Rule;
import org.redisson.api.RScript;
import org.redisson.api.RedissonClient;
import org.redisson.client.codec.LongCodec;

import java.util.Collections;

/**
 *
 * -- 1. rateLimiterKey redis key
 * -- 2. permitsPerSecond  每秒存入的令牌数
 * -- 3. maxBurstSeconds   总共多少秒 （最大存储 maxBurstSeconds 秒生成的令牌）
 * -- 4. needPermits 获取的令牌数
 * -- 5. currentMillis 当前的时间戳
 */
public class SmoothBurstyRateLimiter implements RateLimiter {

    private final RScript rScript;
    // 防止每次的都加载脚本
    private final String sha1;

    public SmoothBurstyRateLimiter(RedissonClient client) {
        this.rScript = client.getScript(LongCodec.INSTANCE);
        String script = LuaScript.getSmoothBurstyRateLimiterScript();
        sha1 = rScript.scriptLoad(script);
    }

    /**
     * 支持预取，不支持下次阻塞，直接返回失败
     */
    @Override
    public Result isAllowed(Rule rule) {
        String rateLimiterKey = rule.getKey();
        int permitsPerSecond = rule.getRate();
        int maxBurstSeconds = rule.getBucketCapacity() / permitsPerSecond;
        int needPermits = rule.getRequestedTokens();
        long currentMillis = System.currentTimeMillis();
        Long r = rScript.evalSha(RScript.Mode.READ_WRITE, sha1,
                RScript.ReturnType.INTEGER, Collections.singletonList(rateLimiterKey),
                permitsPerSecond, maxBurstSeconds, needPermits, currentMillis
        );
        return new Result(r == 0,r);
    }
}
